import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.model_selection import GridSearchCV
import joblib
from utils.log import Logger

logging_obj = Logger(root_path="../",log_name="train")
logger = logging_obj.get_logger()

def init(file_path):
    logger.info('开始读取数据')
    df = pd.read_csv(file_path)
    logger.info('数据读取完毕')
    return df


def preprocess(df):
    logger.info('开始处理数据')
    # drop_columns = ['EducationField','Gender','Education','EmployeeNumber','Over18','PerformanceRating','StandardHours','RelationshipSatisfaction','DistanceFromHome','EmployeeNumber','NumCompaniesWorked','PercentSalaryHike','TrainingTimesLastYear','YearsSinceLastPromotion']
    drop_columns = ['EmployeeNumber','Over18','StandardHours']
    df = df.drop(columns=drop_columns)
    df = pd.get_dummies(df)
    logger.info('数据处理完毕')
    return df


def train(df):
    logger.info('开始训练模型')
    x = df.drop(columns=['Attrition'])
    y = df['Attrition']
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=985,stratify=y)
    transformer = StandardScaler()
    x_train = transformer.fit_transform(x_train)
    x_test = transformer.transform(x_test)
    model = GradientBoostingClassifier(n_estimators=300,max_depth=3,random_state=985,learning_rate=0.05)
    # model = GradientBoostingClassifier()
    # param = {
    #     'learning_rate': [0.01, 0.05, 0.1, 0.2],
    #     'n_estimators': [100,200,300,400,500],
    #     'max_depth': [3,5,7,10,15,20,25],
    #     'random_state':[985]}
    # model = GridSearchCV(model, param_grid=param, cv=5, scoring='roc_auc')
    model.fit(x_train, y_train)
    # logger.info(f'最优参数:{model.best_params_}')
    y_predict = model.predict(x_test)
    y_predict_proba = model.predict_proba(x_test)[:, 1]
    logger.info(f"准确率: {model.score(x_test, y_test)}")
    logger.info(f"AUC分数: {roc_auc_score(y_test, y_predict_proba)}")
    logger.info(f"精确率: {precision_score(y_test, y_predict)}")
    logger.info(f"召回率: {recall_score(y_test, y_predict)}")
    joblib.dump(model, '../model/model.pkl')
    joblib.dump(transformer, '../model/transformer.pkl')
    logger.info('模型训练完毕')




if __name__ == '__main__':
    init_data = init('../data/train.csv')
    df = preprocess(init_data)
    train(df)